#!/usr/bin/env python
import os
import subprocess
import ROOT
ROOT.gROOT.SetBatch()

from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('files', nargs='*')
args = parser.parse_args()

preamble = """
import ROOT
ROOT.gROOT.SetBatch()

def raw_input():
    pass
"""

preamble_rootpy = preamble + """
import rootpy
rootpy.log.basic_config_colorized()
"""

showlabel = False

def walk_files():
    for dirpath, _, filenames in os.walk('.'):
        yield dirpath, filenames

def walk_args():
    for filepath in args.files:
        head, tail = os.path.split(filepath)
        yield head, [tail]

if args.files:
    walk_func = walk_args
else:
    walk_func = walk_files

for dirpath, filenames in walk_func():
    for filename in filenames:
        name, ext = os.path.splitext(filename)
        if ext != '.py':
            continue
        if 'rootpy' in name:
            listingname = 'rootpy'
        else:
            listingname = 'PyROOT'
        content = open(filename, 'r').readlines()
        with open(os.path.join(dirpath, name + '.tex'), 'w') as f:
            f.write('\\begin{footnotesize}\n')
            if showlabel:
                caption = content[0].strip('# ')
                f.write('\\begin{minted}{python}\n'
                        % (caption, listingname))
            else:
                f.write('\\begin{minted}{python}\n')
            if content[0].startswith('"'):
                f.writelines(content)
            else:
                f.writelines(content[1:])
            f.write('\\end{minted}\n')
            f.write('\\end{footnotesize}\n')
        # run the script and write the output
        if os.path.exists(os.path.join(dirpath, name + '_output.lock')):
            continue
        if listingname == 'rootpy':
            content[0:0] = preamble_rootpy
        else:
            content[0:0] = preamble
        tmp_file_name = filename + '.tmp'
        tmp_file = open(tmp_file_name, 'w')
        tmp_file.writelines(content)
        tmp_file.close()
        print "running {0} ...".format(tmp_file_name)
        proc = subprocess.Popen(['python', tmp_file_name],
                stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        output = proc.communicate()[0]
        with open(os.path.join(dirpath, name + '_output.tex'), 'w') as f:
            f.write('\\begin{footnotesize}\n')
            f.write('\\begin{minted}{text}\n')
            f.write(output)
            f.write('\\end{minted}\n')
            f.write('\\end{footnotesize}\n')
        os.unlink(tmp_file_name)
